import datetime
import os
import socket
import sys
import threading
import time

import numpy as np
import torch
from PySide6.QtCore import QFile
from PySide6.QtGui import QPixmap
from PySide6.QtWidgets import (
    QApplication,
    QLabel,
    QMainWindow,
    QHeaderView,
    QTableWidgetItem,
    QAbstractItemView,
)
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from table import Base, Data, Device, TimeError
from ui_mainwindow import Ui_MainWindow
from util import getMacByIp


def getIp():
    """
    查询本机的局域网ip地址
    :return: ip
    """
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
    finally:
        s.close()
    return ip


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.imgUpdate(
            r"/run/media/kearney/a/CAU/42course/毕设/数据集/NEU-DET/IMAGES/patches_253.jpg"
        )
        self.numSum = 0
        self.numError = 0
        self.initTableHeader()
        self.initDataBase()
        self.initServer()

    def initTableHeader(self):
        """设置表头"""
        self.ui.tableLogin.setEditTriggers(QAbstractItemView.NoEditTriggers)  # 设置表格不可编辑
        self.ui.tableLogin.setColumnCount(3)  # 3 列
        self.ui.tableLogin.setHorizontalHeaderLabels(["Mac", "IP", "Time"])  # 设置表头
        self.ui.tableLogin.horizontalHeader().setSectionResizeMode(
            QHeaderView.Stretch
        )  # 宽度自适应 Stretch：宽度适应界面自动伸缩；ResizeToContents：宽度根据内容自适应
        self.ui.tabData.setColumnCount(4)
        self.ui.tabData.setHorizontalHeaderLabels(
            ["Mac", "Time", "Error Number", "Photo Number"]
        )
        self.ui.tabData.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)

        self.ui.tabTime.setColumnCount(4)
        self.ui.tabTime.setHorizontalHeaderLabels(["Mac", "Time", "Class", "Number"])
        self.ui.tabTime.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)

    def initDataBase(self):
        """初始化数据库，向界面上填入数据"""
        engine = create_engine(
            "sqlite:///maxipy.db?check_same_thread=False",
            echo=False,
            future=True,
            encoding="utf-8",
        )  # 数据库连接
        Base.metadata.create_all(engine)  # 建立表，已有表就不会重新建立了
        DbSession = sessionmaker(bind=engine)  # 创建session
        self.session = DbSession()  # 建立会话

        # 向界面上的 table 插入数据
        res = self.session.query(Device).all()
        self.ui.tableLogin.setRowCount(len(res))  # 设置行数
        for i in range(len(res)):  # 循环给每一个单元格赋值
            self.ui.tableLogin.setItem(i, 0, QTableWidgetItem(str(res[i].mac)))
            self.ui.tableLogin.setItem(i, 1, QTableWidgetItem(str(res[i].ip)))
            self.ui.tableLogin.setItem(i, 2, QTableWidgetItem(str(res[i].time)))

        res = self.session.query(Data).all()
        self.ui.tabData.setRowCount(len(res))
        for i in range(len(res)):
            self.ui.tabData.setItem(i, 0, QTableWidgetItem(str(res[i].mac)))
            self.ui.tabData.setItem(i, 1, QTableWidgetItem(str(res[i].time)))
            self.ui.tabData.setItem(i, 2, QTableWidgetItem(str(res[i].errorNum)))
            self.ui.tabData.setItem(i, 3, QTableWidgetItem(str(res[i].photoNum)))

        res = self.session.query(TimeError).all()
        self.ui.tabTime.setRowCount(len(res))
        for i in range(len(res)):
            self.ui.tabTime.setItem(i, 0, QTableWidgetItem(str(res[i].mac)))
            self.ui.tabTime.setItem(i, 1, QTableWidgetItem(str(res[i].time)))
            self.ui.tabTime.setItem(i, 2, QTableWidgetItem(str(res[i].objcls)))
            self.ui.tabTime.setItem(i, 3, QTableWidgetItem(str(res[i].num)))

    def initServer(self, localPort=8888):
        """初始化图传服务

        localPort：服务端口"""
        # 在底部状态栏显示本机局域网 ip
        self.showTimingTemMessage("Init Internet")
        self.localIp = getIp()
        self.showPermanentMessage("IP: %s Port: %s" % (self.localIp, localPort))

        self.localPort = localPort
        self.sk = socket.socket()
        self.sk.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sk.bind((self.localIp, self.localPort))
        self.sk.listen(50)
        self.showTimingTemMessage("socket ok, wait for client")

        self.thread = threading.Thread(target=self.server, args=())
        self.thread.daemon = True
        self.thread.start()

        self.showTimingTemMessage("Init Internet success")

    def server(self):
        """监听线程，有客户端连上就给客户开新线程专线服务"""
        while True:
            conn, addr = self.sk.accept()
            print("hello client,ip:", addr)
            self.ui.status.setText("Connected")  # 更新界面上显示的 Status
            self.ui.ip.setText(str(addr[0]))  # 更新界面上显示的 IP
            mac = getMacByIp(addr[0])
            self.ui.mac.setText(mac)  # 更新界面上显示的 MAC,测试发现有点延迟
            # 设备上线，录入数据库
            DeviceLogin = Device(mac, str(addr[0]), datetime.datetime.now())
            self.session.add(DeviceLogin)
            self.session.commit()
            self.session.commit()

            t = threading.Thread(target=self.receiveThread, args=(conn,))
            t.daemon = True
            t.start()

    def receiveThread(self, conn):
        """
        图像接收线程，接收到图像后调用模型进行预测，并将结果返回给客户端
        """
        # conn.settimeout(10)
        conn_end = False
        pack_size = 100000  # 10w
        print("begin to receive")
        while True:
            img = b""
            client_data = None
            errTime = 0  # 错误计数
            if conn_end:
                break
            while True:
                try:
                    client_data = conn.recv(2048)
                    # print("rcv:", len(client_data))
                    time.sleep(0.015)
                except Exception as e:
                    print(e)
                    break
                if not client_data:
                    # print("received null")
                    errTime += 1
                    if errTime > 299:  # 多次收到空包，认定为结束
                        conn_end = True
                        break
                    continue
                else:
                    img += client_data
                if client_data.endswith(b"\xFF\xD9"):  # jpg 尾部检验
                    print("image end dect. recive end, pic len:", len(img))
                    f = open("tmp.jpg", "wb")  # 写入图片
                    f.write(img)
                    f.close()
                    errTime = 0  # 错误次数清零
                    self.numSum += 1
                    self.numError += 1
                    self.numUpdate(self.numError, self.numSum)
                    self.imgUpdate("tmp.jpg")
                    res, NumC = getPre(MODEL, "tmp.jpg", 200)  # 使用模型进行目标检测
                    self.imgUpdate("tmp.jpg")
                    conn.send(res.encode("utf-8"))
                    # 录入数据库
                    ## Data photoNum+=1
                    if (
                        self.session.query(Data)
                        .filter(
                            Data.mac == self.ui.mac.text().strip(),
                            Data.time
                            == datetime.datetime.now().strftime("%Y-%m-%d %H"),
                        )
                        .count()
                        == 0
                    ):  # 查找有没有
                        init = Data(
                            self.ui.mac.text().strip(),
                            datetime.datetime.now().strftime("%Y-%m-%d %H"),
                            (1 if len(NumC) > 0 else 0),
                            1,
                        )
                        self.session.add(init)
                        self.session.commit()
                    else:
                        temp = (
                            self.session.query(Data)
                            .filter(
                                Data.mac == self.ui.mac.text().strip(),
                                Data.time
                                == datetime.datetime.now().strftime("%Y-%m-%d %H"),
                            )
                            .one()
                        )
                        temp.photoNum += 1
                        temp.errorNum += 1 if len(NumC) > 0 else 0
                        self.session.commit()
                    ## TimeError
                    for i in NumC.index:  # i：index， NumC[i]：num
                        try:
                            temp = (
                                self.session.query(TimeError)
                                .filter(
                                    TimeError.mac == self.ui.mac.text().strip(),
                                    TimeError.time
                                    == datetime.datetime.now().strftime("%Y-%m-%d %H"),
                                    TimeError.objcls == i,
                                )
                                .one()
                            )
                            # 更新数据
                            temp.num += int(NumC[i])
                            self.session.commit()
                        except Exception as e:
                            print("TimeError: ", i, NumC[i], e)
                            # No row was found when one was required 说明不存在这一行数据，则要进行插入
                            add = TimeError(
                                self.ui.mac.text().strip(),
                                datetime.datetime.now().strftime("%Y-%m-%d %H"),
                                i,
                                int(NumC[i]),
                            )
                            self.session.add(add)
                            self.session.commit()
                    break
                if len(img) > pack_size:  # 数据异常大
                    print("too big, unabnomal, clear")
                    conn_end = True
                    break
        conn.close()
        print("receive thread end")

    def imgUpdate(self, imgPath):
        """更新界面显示的图像（Error,Sum）"""
        self.ui.imgPre.setPixmap(QPixmap(imgPath))

    def numUpdate(self, numError, numSum):
        """更新界面显示的数据（Error,Sum）"""
        self.ui.numError.setText(str(numError))
        self.ui.numSum.setText(str(numSum))

    def showTimingTemMessage(self, msg="temp message", time=2000):
        """显示临时消息 定时ms后消失"""
        self.ui.statusbar.showMessage(msg, time)

    def showTempMessage(self, msg="temp message"):
        """显示临时消息"""
        self.ui.statusbar.showMessage(msg)

    def showNormalMessage(self, msg="message"):
        """显示普通消息"""
        normalMessage = QLabel(msg)
        normalMessage.setObjectName("normalMessage")
        self.ui.statusbar.addWidget(normalMessage)

    def showPermanentMessage(self, msg="message"):
        """持久显示消息"""
        permanentMessage = QLabel(msg)
        permanentMessage.setObjectName("permanentMessage")
        self.ui.statusbar.addPermanentWidget(permanentMessage)

    def removeNormalMessage(self):
        """移除普通消息"""
        self.ui.statusbar.removeWidget(
            self.ui.statusbar.findChild(QLabel, "normalMessage")
        )
        self.ui.statusbar.removeWidget(
            self.ui.statusbar.findChild(QPushButton, "normalBtn")
        )

    def removePermanentMessage(self):
        """移除持久消息"""
        self.ui.statusbar.removeWidget(
            self.ui.statusbar.findChild(QLabel, "permanentMessage")
        )
        self.ui.statusbar.removeWidget(
            self.ui.statusbar.findChild(QPushButton, "permanentBtn")
        )


def xyxy2xywh(x):
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = x[:, 0]  # x1=top-left.x
    y[:, 1] = x[:, 1]  # y1=top-left.y
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y.int()


def getPre(model, img, imgSize):
    """
    参数：
    model：yolo5模型
    img：图像路径
    imgSize：方形图像大小，如 200
    """
    results = model(img, size=imgSize)  # 指定图像大小，不指定的效果很差
    results.save(os.getcwd())  # 将预测结果画到图片里，并保存到当前工作目录
    pre = results.pandas().xyxy[0].sort_values("xmin")
    tmp = pre.iloc[:, 0:4]
    NumC = pre["name"].value_counts()  # 错误类型和数量
    temp = np.array(tmp).reshape(-1, 4)
    xywh = (xyxy2xywh(torch.tensor(temp))).tolist()  # normalized xywh
    lines = ""
    for i in range(len(pre)):
        tmpres = pre.loc[i, "name"], " ".join(map(str, xywh[i]))
        lines += (
            (str(tmpres) + ("\n" if i < len(pre) - 1 else ""))
            .replace(",", "")
            .replace("'", "")
            .replace("(", "")
            .replace(")", "")
        )
    return lines, NumC


if __name__ == "__main__":
    # Model
    MODEL = torch.hub.load(
        "../yolov5",
        "custom",
        path="../yolov5/runs/train/exp11/weights/best",
        source="local",
    )
    app = QApplication(sys.argv)

    window = MainWindow()
    window.show()

    sys.exit(app.exec())
